#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <netinet/in.h>
#include <memory.h>
#include "list_dir.h"
#include "RetrievePcap.h"

#define equal(a,b) ((a->BasicInfo).Ip.Dst==(b->BasicInfo).Ip.Dst)&&((a->BasicInfo).Ip.Src==(b->BasicInfo).Ip.Src)&&((a->BasicInfo).Port.Src==(b->BasicInfo).Port.Src)&&((a->BasicInfo).Port.Dst==(b->BasicInfo).Port.Dst)&&((a->BasicInfo).Protocol==(b->BasicInfo).Protocol)
#define WIDTH 34

//construct a new flow with the new packet not belong to any existent flow
void process_flow(struct flow *FlowHead,struct flow **FlowCurrent,struct flow* FlowTemp,struct packet *newPacket,char* protocol,struct statistic* Statistic)
{
    FILE* fp;
    int i;
    for(FlowTemp=FlowHead->next;FlowTemp!=NULL&&!(equal(FlowTemp,newPacket));FlowTemp=FlowTemp->next);
    if(FlowTemp!=NULL)
    {
        if(FlowTemp->count<5)
            FlowTemp->PKG_SIZE[FlowTemp->count++]=newPacket->captureLen;
        if(FlowTemp->count==5)
        {
            if((fp=fopen("TEMP","a+"))==NULL) 
            {
                printf("Creating TEMP file...\n");   
                if((fp=fopen("TEMP","w"))==NULL) 
                {
                    printf("Error creating TEMP file...\n");   
                    exit(1);
                }
            }
            if((FlowTemp->BasicInfo).Protocol==IPPROTO_TCP)
                fprintf(fp,"TCP,");
            else
                fprintf(fp,"UDP,");
            fprintf(fp,"%d,",(FlowTemp->BasicInfo).Port.Src);
            fprintf(fp,"%d,",(FlowTemp->BasicInfo).Port.Dst);
            for(i=0;i<5;i++)
                fprintf(fp,"%d,",FlowTemp->PKG_SIZE[i]);
            fprintf(fp,"%s\n",protocol); 
            fclose(fp);
            FlowTemp->count=6;    
            Statistic->TotOfAail++;
            Statistic->SingleTotOfAail++;
            //  printf("-\n");
        }
    }
    else
    {
        FlowTemp=(struct flow*)malloc(sizeof(struct flow));
        memset(FlowTemp,0,sizeof(struct flow));
        FlowTemp->PKG_SIZE[FlowTemp->count++]=newPacket->captureLen;
        (FlowTemp->BasicInfo).Ip.Src=(newPacket->BasicInfo).Ip.Src;
        (FlowTemp->BasicInfo).Ip.Dst=(newPacket->BasicInfo).Ip.Dst;
        (FlowTemp->BasicInfo).Port.Src=(newPacket->BasicInfo).Port.Src;
        (FlowTemp->BasicInfo).Port.Dst=(newPacket->BasicInfo).Port.Dst;
        (FlowTemp->BasicInfo).Protocol=(newPacket->BasicInfo).Protocol;
        FlowTemp->next=NULL;
        (*FlowCurrent)->next=FlowTemp;
        (*FlowCurrent)=FlowTemp;
        Statistic->TotOfFlow++;
        Statistic->SingleTotOfFlow++;
    }
}
// process a packet in pcap file, record info to struct packet
int process_packet(FILE* fp,struct packet *newPacket)
{
    unsigned int captureLen,ipHeaderLen;
    unsigned long ipSrc,ipDst;
    unsigned short portSrc,portDst;
    int i;
    skip(fp, 8);

    //captureLen
    process_4B_uint(fp, &captureLen);
    newPacket->captureLen=captureLen;
    //printf("CaptureLen=%d\n",captureLen);
    skip(fp, 17); //ethernet header
    // printf("IP=%d,fgetc(fp)=%d\n",IPPROTO_IP,(unsigned char)fgetc(fp));
    if(fgetc(fp)!=IPPROTO_IP)
    {
        skip(fp, captureLen-14); //packet is second layer protocol,we skip it  
        return 0;
    }
    ipHeaderLen=(unsigned int)(fgetc(fp)& 0x0f)*4;
    newPacket->ipHeaderLen=ipHeaderLen;
    skip(fp, 8);
    (newPacket->BasicInfo).Protocol=(unsigned char)fgetc(fp);
    skip(fp,2); //checksum

    process_ip(fp, &ipSrc);
    process_ip(fp, &ipDst);
    (newPacket->BasicInfo).Ip.Src=ipSrc;
    (newPacket->BasicInfo).Ip.Dst=ipDst;
    skip(fp, ipHeaderLen-20);
    // printf("CaptureLen:%d,IpHeaderLen:%d,Protocol:%d\n",captureLen,ipHeaderLen,(newPacket->BasicInfo).Protocol);
    //   scanf("%d",&i);
    if((newPacket->BasicInfo).Protocol==IPPROTO_TCP||(newPacket->BasicInfo).Protocol==IPPROTO_UDP) //it is a tcp or udp packet
    {
        process_port(fp, &portSrc);
        process_port(fp, &portDst);
        (newPacket->BasicInfo).Port.Src=portSrc;
        (newPacket->BasicInfo).Port.Dst=portDst;
        skip(fp, captureLen-18-ipHeaderLen);
    }
    else 
    {
        skip(fp, captureLen-14-ipHeaderLen); //packet is neither tcp nor udp, we ignore it 
        return 0;
    }
    return 1;
}
/*判断文件是否是个目录文件，是目录返回 1 ，不是目录返回 0*/
int testdir(char *path)
{
    struct stat buf;
    if(lstat(path,&buf)<0)
    {
        return 0;
    }
    if(S_ISDIR(buf.st_mode))
    {
        return 1; //directory
    }
    return 0;
}
void free_flow_list(struct flow* FlowHead,int RemoveHead)
{
    struct flow  *FlowCurrent,*FlowTemp;
    if(RemoveHead)
        for(FlowCurrent=FlowHead;FlowCurrent!=NULL;)
        {
            FlowTemp=FlowCurrent->next;
            free(FlowCurrent);
            FlowCurrent=FlowTemp;
        }
    else
        for(FlowCurrent=FlowHead->next;FlowCurrent!=NULL;)
        {
            FlowTemp=FlowCurrent->next;
            free(FlowCurrent);
            FlowCurrent=FlowTemp;
        }
}
int IsPcapFile(char* filename)
{
    FILE* fp;
    if((fp=fopen(filename,"rb"))==NULL) 
    {
        fclose(fp);
        printf("File %s Open Error...\n",filename);
        return 0;
    }
    if(fgetc(fp)==212&&fgetc(fp)==195&&fgetc(fp)==178&&fgetc(fp)==161
            &&fgetc(fp)==2&&fgetc(fp)==0
            &&fgetc(fp)==4&&fgetc(fp)==0)
    {
        fclose(fp);
        return 1; 
    }
    fclose(fp);
    return 0;
}
/*遍历目录*/
int directory(char *path)
{
    DIR *db;              /*保存 打开目录类型文件信息的 结构体*/
    char filename[FILE_NAME],file_name[FILE_NAME];        /*文件名*/
    struct stat file_info;         /*保存 普通文件的属性信息的 结构体*/
    struct type *head,*current,*temp;
    char *delim="_",*protocol;
    FILE *fp,*fp2;
    struct dirent *dir_info;        /*保存 目录类型文件属性信息的 结构体*/
    struct packet* newPacket; 
    struct flow  *FlowHead,*FlowCurrent,*FlowTemp;
    struct statistic Statistic;
    struct timeval start,end;
    char FilenameAndSize[40];
    int i;
    unsigned long timeuse; 

    memset(&Statistic,0,sizeof(struct statistic));
    newPacket=(struct packet*)malloc(sizeof(struct packet));

    FlowHead=(struct flow*)malloc(sizeof(struct flow));
    memset( FlowHead,0,sizeof(struct flow));   /*将文件信息结构体置 0 */

    head=(struct type*)malloc(sizeof(struct type));
    memset( head,0,sizeof(struct type));   /*将文件信息结构体置 0 */

    head->next=NULL;
    current=head;


    db=opendir(path);
    if(db==NULL){
        printf("Open dir error !\n");
        return 0;
    }
    memset(filename,0,FILE_NAME);

    while ((dir_info=readdir(db)))
    {
        if((strcmp(dir_info->d_name,".")==0)||(strcmp(dir_info->d_name,"..")==0))
            continue;
        else
        {
            sprintf(filename,"%s/%s",path,dir_info->d_name);
            if(!testdir(filename)&&IsPcapFile(filename))
            {
                memset( &file_info,0,sizeof(file_info));   /*将文件信息结构体置 0 */
                lstat( filename, &file_info );      /*通过文件路径，把文件信息结构体填满*/
                gettimeofday(&start,NULL);
                if((fp=fopen(filename,"rb"))==NULL)
                {
                    printf("> File doesn't exist!\n");
                    exit(1);
                }
                skip(fp,24);
                Statistic.SingleTotOfFlow=0;
                Statistic.SingleTotOfAail=0;
                if(file_info.st_size<1024)
                    sprintf(FilenameAndSize,"%-*sSize:%8dB",WIDTH,
                            file_name,file_info.st_size);
                else
                    if(file_info.st_size/1024<1024)
                        sprintf(FilenameAndSize,"%-*sSize:%7.3fKB",WIDTH,
                                file_name,((double)(file_info.st_size))/1024);
                    else
                        if(file_info.st_size/1024/1024<1024)
                            sprintf(FilenameAndSize,"%-*sSize:%7.3fMB",WIDTH,
                                    file_name,((double)(file_info.st_size))/1024/1024);
                        else
                            if(file_info.st_size/1024/1024/1024<1024)
                                sprintf(FilenameAndSize,"%-20sSize:%7.3fGB",WIDTH,
                                        file_name,((double)(file_info.st_size))/1024/1024/1024);
                FlowHead->next=NULL;
                FlowCurrent=FlowHead;
                while(!feof(fp))
                {
                    memset(newPacket,0,sizeof(struct packet));
                    Statistic.TotOfPkg++;
                    if(process_packet(fp,newPacket,protocol))
                    {
                        process_flow(FlowHead,&FlowCurrent,FlowTemp,newPacket,protocol,&Statistic);
                        for(temp=head->next;temp!=NULL&&strcmp(temp->ProtocolName,protocol);temp=temp->next);
                        if(temp==NULL) 
                        {
                            temp=(struct type*)malloc(sizeof(struct type));
                            memset(temp,0,sizeof(struct type));   /*将文件信息结构体置 0 */
                            strcpy(temp->ProtocolName,protocol);
                            temp->next=NULL;
                            current->next=temp;
                            current=temp;
                        }

                    }
                }
               // free_flow_list(FlowHead,0);
                fclose(fp);
                gettimeofday(&end,NULL);
                timeuse=1000000*(end.tv_sec-start.tv_sec)+end.tv_usec-start.tv_usec;
                Statistic.TotOfSize+=file_info.st_size;
                printf("%s\tTotal number of flows:%5d\tTotal number of useful flows:%5d\tRun time is %.3fs\n",FilenameAndSize,Statistic.SingleTotOfFlow,Statistic.SingleTotOfAail,((double)(timeuse))/1000000);
                fflush(stdout);
            }
        }
        memset(filename,0,64);
    }
    if(Statistic.TotOfSize<1024)
        printf("%*sTotal size:%8dB\tTotal number of flows:%5d\tTotal number of useful flows:%5d\t",WIDTH-6,
                "",Statistic.TotOfSize,Statistic.TotOfFlow,Statistic.TotOfAail);
    else
        if(Statistic.TotOfSize/1024<1024)
            printf("%*sTotal Size:%7.3fKB\tTotal number of flows:%5d\tTotal number of useful flows:%5d\t",WIDTH-6,
                    "",((double)(Statistic.TotOfSize))/1024,Statistic.TotOfFlow,Statistic.TotOfAail);
        else
            if(Statistic.TotOfSize/1024/1024<1024)
                printf("%*sTotal Size:%7.3fMB\tTotal number of flows:%5d\tTotal number of useful flows:%5d\t",WIDTH-6,
                        "",((double)(Statistic.TotOfSize))/1024/1024,Statistic.TotOfFlow,Statistic.TotOfAail);
            else
                if(Statistic.TotOfSize/1024/1024/1024<1024)
                    printf("%*sTotal Size:%7.3fGB\tTotal number of flows:%5d\tTotal number of useful flows:%5d\t",WIDTH-6,
                            "",((double)(Statistic.TotOfSize))/1024/1024/1024,Statistic.TotOfFlow,Statistic.TotOfAail);
    closedir(db);
    if(head->next==NULL)
    {
        printf("There is no pcap file in that directory or file name is not right...\n");
        exit(1);
    }
    current=head->next;
    if((fp=fopen("Result.arff","w"))==NULL)
    {
        printf("> File doesn't exist!\n");
        exit(1);
    }
    fprintf(fp,"@relation Flow_Classification\n\n");
    fprintf(fp,"@attribute TCP_OR_UDP {TCP,UDP}\n");
    fprintf(fp,"@attribute SRC_PORT numeric\n");
    fprintf(fp,"@attribute DST_PORT numeric\n");
    fprintf(fp,"@attribute PKG_SIZE_1 numeric\n");
    fprintf(fp,"@attribute PKG_SIZE_2 numeric\n");
    fprintf(fp,"@attribute PKG_SIZE_3 numeric\n");
    fprintf(fp,"@attribute PKG_SIZE_4 numeric\n");
    fprintf(fp,"@attribute PKG_SIZE_5 numeric\n");
    fprintf(fp,"@attribute PROTOCOL {%s",current->ProtocolName);
    for(current=current->next;current!=NULL;current=current->next)
        fprintf(fp,",%s",current->ProtocolName);
    fprintf(fp,"}\n\n");
    fprintf(fp,"@data\n\n");
    fclose(fp);
    system("cat TEMP >> Result.arff && rm -f TEMP");

    for(current=head;current!=NULL;)
    {
        temp=current->next;
        free(current);
        current=temp;
    }
    free_flow_list(FlowHead,1);
    free(newPacket);
    return 0;
}

//program starts here
int main(int argc, char *argv[])
{
    struct timeval start,end;
    unsigned long timeuse;
    system("rm -f TEMP");
    gettimeofday(&start,NULL);
    if(argc<2)
    {
        printf("pcap_retrieve.exe [The path of directory]\n");
        return 0;
    }
    directory(argv[1]);
    gettimeofday(&end,NULL);
    timeuse=1000000*(end.tv_sec-start.tv_sec)+end.tv_usec-start.tv_usec;
    printf("Run time is %.3fs\n",((double)(timeuse))/1000000);
    fflush(stdout);
    system("java weka.classifiers.trees.J48 -C 0.25 -M 2 -x 10 -t \"./Result.arff\"");
    return 0;
}
